// Copyright (c) 2023, AgiBot Inc.
// All rights reserved.

#include <string>
#include <unordered_map>
#include <vector>

#include <google/protobuf/compiler/code_generator.h>
#include <google/protobuf/compiler/plugin.h>
#include <google/protobuf/compiler/plugin.pb.h>
#include <google/protobuf/descriptor.h>

std::string& ReplaceString(std::string& source, const std::string& replace_what,
                           const std::string& replace_with_what) {
  std::string::size_type pos = 0;
  while (true) {
    pos = source.find(replace_what, pos);
    if (pos == std::string::npos) break;
    source.replace(pos, replace_what.size(), replace_with_what);
    pos += replace_with_what.size();
  }

  return source;
}

std::vector<std::string> SplitToVec(const std::string& source,
                                    const std::string& sep, bool clear = true) {
  std::vector<std::string> result;
  if (source.empty() || sep.empty()) return result;

  size_t pos_end = 0, pos_start = 0;
  do {
    pos_end = source.find(sep, pos_start);
    if (pos_end == std::string::npos) pos_end = source.length();

    const std::string& sub_str = source.substr(pos_start, pos_end - pos_start);
    if (!(clear && sub_str.empty())) {
      result.emplace_back(sub_str);
    }

    pos_start = pos_end + sep.size();
  } while (pos_end < source.length());

  return result;
}

class AimRTCodeGenerator final : public google::protobuf::compiler::CodeGenerator {
 public:
  AimRTCodeGenerator() = default;
  ~AimRTCodeGenerator() override = default;

  uint64_t GetSupportedFeatures() const override {
    return FEATURE_PROTO3_OPTIONAL;
  }

  constexpr static std::string_view kTHfile = R"str(/**
 * @file {{file_name}}.aimrt_rpc.pb.h
 * @brief This file was generated by protoc-gen-aimrt_rpc which is a self-defined pb compiler plugin, do not edit it!!!
 */
#pragma once

#include <future>

#include "aimrt_module_cpp_interface/rpc/rpc_handle.h"
#include "aimrt_module_cpp_interface/rpc/rpc_status.h"
#include "aimrt_module_cpp_interface/util/version.h"

#include "aimrt_module_cpp_interface/co/task.h"

#include "aimrt_module_cpp_interface.h"

#include "{{file_name}}.pb.h"

static_assert({{cur_gencode_version}} <= AIMRT_RUNTIME_VERSION_INT,
              "AIMRT_RUNTIME_VERSION is older than generated code version {{cur_gencode_version_str}}");
static_assert(AIMRT_MIN_GENCODE_VERSION_INT <= {{cur_gencode_version}},
              "AIMRT_MIN_GENCODE_VERSION is greater than generated code version {{cur_gencode_version_str}}");


{{namespace_begin}}
{{for service begin}}
class {{service_name}}SyncService : public aimrt::rpc::ServiceBase {
 public:
  {{service_name}}SyncService();
  ~{{service_name}}SyncService() override = default;
{{for method begin}}
  virtual aimrt::rpc::Status {{rpc_func_name}}(
      aimrt::rpc::ContextRef ctx_ref,
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp) {
    return aimrt::rpc::Status(AIMRT_RPC_STATUS_SVR_NOT_IMPLEMENTED);
  }
{{method end}}
};
{{service end}}

{{for service begin}}
class {{service_name}}AsyncService : public aimrt::rpc::ServiceBase {
 public:
  {{service_name}}AsyncService();
  ~{{service_name}}AsyncService() override = default;
{{for method begin}}
  virtual void {{rpc_func_name}}(
      aimrt::rpc::ContextRef ctx_ref,
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp,
      std::function<void(aimrt::rpc::Status)>&& callback) {
    callback(aimrt::rpc::Status(AIMRT_RPC_STATUS_SVR_NOT_IMPLEMENTED));
  }
{{method end}}
};
{{service end}}

{{for service begin}}
class {{service_name}}CoService : public aimrt::rpc::CoServiceBase {
 public:
  {{service_name}}CoService();
  ~{{service_name}}CoService() override = default;
{{for method begin}}
  virtual aimrt::co::Task<aimrt::rpc::Status> {{rpc_func_name}}(
      aimrt::rpc::ContextRef ctx_ref,
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp) {
    co_return aimrt::rpc::Status(AIMRT_RPC_STATUS_SVR_NOT_IMPLEMENTED);
  }
{{method end}}
};

{{service end}}

{{for service begin}}
bool Register{{service_name}}ClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name);
bool Register{{service_name}}ClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref);
{{service end}}

{{for service begin}}
class {{service_name}}SyncProxy : public aimrt::rpc::SyncProxyBase {
 public:
  explicit {{service_name}}SyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref);
  explicit {{service_name}}SyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name);
  ~{{service_name}}SyncProxy() = default;

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref);
  }

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref, service_name);
  }
{{for method begin}}
  aimrt::rpc::Status {{rpc_func_name}}(
      aimrt::rpc::ContextRef ctx_ref,
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp);

  aimrt::rpc::Status {{rpc_func_name}}(
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp) {
    return {{rpc_func_name}}(aimrt::rpc::ContextRef(), req, rsp);
  }
{{method end}}
};
{{service end}}

{{for service begin}}
class {{service_name}}AsyncProxy : public aimrt::rpc::AsyncProxyBase {
 public:
  explicit {{service_name}}AsyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref);
  explicit {{service_name}}AsyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name);
  ~{{service_name}}AsyncProxy() = default;

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref);
  }

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref, service_name);
  }
{{for method begin}}
  void {{rpc_func_name}}(
      aimrt::rpc::ContextRef ctx_ref,
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp,
      std::function<void(aimrt::rpc::Status)>&& callback);

  void {{rpc_func_name}}(
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp,
      std::function<void(aimrt::rpc::Status)>&& callback) {
    {{rpc_func_name}}(aimrt::rpc::ContextRef(), req, rsp, std::move(callback));
  }
{{method end}}
};
{{service end}}


{{for service begin}}
class {{service_name}}FutureProxy : public aimrt::rpc::FutureProxyBase {
 public:
  explicit {{service_name}}FutureProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref);
  explicit {{service_name}}FutureProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name);
  ~{{service_name}}FutureProxy() = default;

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref);
  }

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref, service_name);
  }
{{for method begin}}
  std::future<aimrt::rpc::Status> {{rpc_func_name}}(
      aimrt::rpc::ContextRef ctx_ref,
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp);

  std::future<aimrt::rpc::Status> {{rpc_func_name}}(
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp) {
    return {{rpc_func_name}}(aimrt::rpc::ContextRef(), req, rsp);
  }
{{method end}}
};
{{service end}}

{{for service begin}}
class {{service_name}}CoProxy : public aimrt::rpc::CoProxyBase {
 public:
  explicit {{service_name}}CoProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref);
  explicit {{service_name}}CoProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name);
  ~{{service_name}}CoProxy() = default;

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref);
  }

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref, service_name);
  }
{{for method begin}}
  aimrt::co::Task<aimrt::rpc::Status> {{rpc_func_name}}(
      aimrt::rpc::ContextRef ctx_ref,
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp);

  aimrt::co::Task<aimrt::rpc::Status> {{rpc_func_name}}(
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp) {
    return {{rpc_func_name}}(aimrt::rpc::ContextRef(), req, rsp);
  }
{{method end}}
};
{{service end}}

{{for service begin}}
class {{service_name}}CoClient {
 public:
  /**
   * @brief Initialize the service resource based on AimRT Context.
   */
  void Init(const std::shared_ptr<::aimrt::context::Context> & ctx_ptr, std::string service_name = "");

 public:
{{for method begin}}
  ::aimrt::context::res::Client<{{rpc_req_name}}, {{rpc_rsp_name}}> {{rpc_func_name}};

{{method end}}

 public:
  static std::vector<std::string> GetMethodNames(std::string service_name = {});

};
{{service end}}

{{for service begin}}
// Forward declaration
class {{service_name}}AimRTImpl;

class {{service_name}}CoServer : public std::enable_shared_from_this<{{service_name}}CoServer> {
 public:
  void Init(const std::shared_ptr<::aimrt::context::Context> & ctx_ptr, std::string service_name = "");

 public:
{{for method begin}}
  ::aimrt::context::res::Server<{{rpc_req_name}}, {{rpc_rsp_name}}> {{rpc_func_name}};

{{method end}}

 public:
  static std::vector<std::string> GetMethodNames(std::string service_name = {});

 private:
  std::weak_ptr<::aimrt::context::Context> ctx_weak_ptr_;
  std::shared_ptr<{{service_name}}AimRTImpl> impl_;
};

class {{service_name}}AimRTImpl final : public {{service_name}}CoService {
 public:
  explicit {{service_name}}AimRTImpl(
      std::weak_ptr<::aimrt::context::Context> weak_ctx_ptr,
      std::shared_ptr<{{service_name}}CoServer> res);

{{for method begin}}
  aimrt::co::Task<aimrt::rpc::Status> {{rpc_func_name}}(
      aimrt::rpc::ContextRef ctx_ref,
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp) override;

{{method end}}

 private:
  std::weak_ptr<::aimrt::context::Context> weak_ctx_ptr_;
  std::shared_ptr<{{service_name}}CoServer> res_;
};
{{service end}}

{{namespace_end}}
)str";

  constexpr static std::string_view kTCcfile = R"str(/**
 * @file {{file_name}}.aimrt_rpc.pb.cc
 * @brief This file was generated by protoc-gen-aimrt_rpc which is a self-defined pb compiler plugin, do not edit it!!!
 */

#include "{{file_name}}.aimrt_rpc.pb.h"

#include "aimrt_module_cpp_interface/co/inline_scheduler.h"
#include "aimrt_module_cpp_interface/co/on.h"
#include "aimrt_module_cpp_interface/co/start_detached.h"
#include "aimrt_module_cpp_interface/co/then.h"
#include "aimrt_module_protobuf_interface/util/protobuf_type_support.h"

#include "aimrt_module_cpp_interface.h"

{{namespace_begin}}
static constexpr std::string_view kRpcType = "pb";

{{for service begin}}
static constexpr std::string_view k{{service_name}}Name = "{{package_name}}.{{service_name}}";
{{service end}}

{{for service begin}}
{{service_name}}SyncService::{{service_name}}SyncService() : aimrt::rpc::ServiceBase(kRpcType, k{{service_name}}Name) {
{{for method begin}}
  {
    aimrt::rpc::ServiceFunc service_callback(
        [this](const aimrt_rpc_context_base_t* ctx, const void* req, void* rsp, aimrt_function_base_t* result_callback_ptr) {
          aimrt::rpc::ServiceCallback result_callback(result_callback_ptr);

          aimrt::rpc::ContextRef ctx_ref(ctx);

          auto status = {{rpc_func_name}}(
              ctx_ref,
              *static_cast<const {{rpc_req_name}}*>(req),
              *static_cast<{{rpc_rsp_name}}*>(rsp));

          result_callback(status.Code());
        });
    RegisterServiceFunc(
        "{{rpc_func_name}}",
        nullptr,
        aimrt::GetProtobufMessageTypeSupport<{{rpc_req_name}}>(),
        aimrt::GetProtobufMessageTypeSupport<{{rpc_rsp_name}}>(),
        std::move(service_callback));
  }
{{method end}}
}
{{service end}}

{{for service begin}}
{{service_name}}AsyncService::{{service_name}}AsyncService() : aimrt::rpc::ServiceBase(kRpcType, k{{service_name}}Name) {
{{for method begin}}
  {
    aimrt::rpc::ServiceFunc service_callback(
        [this](const aimrt_rpc_context_base_t* ctx, const void* req, void* rsp, aimrt_function_base_t* result_callback_ptr) {
          auto result_callback_func_ptr = std::make_shared<aimrt::rpc::ServiceCallback>(result_callback_ptr);

          aimrt::rpc::ContextRef ctx_ref(ctx);

          {{rpc_func_name}}(
              ctx_ref,
              *static_cast<const {{rpc_req_name}}*>(req),
              *static_cast<{{rpc_rsp_name}}*>(rsp),
              [result_callback_func_ptr{std::move(result_callback_func_ptr)}](aimrt::rpc::Status status) {
                (*result_callback_func_ptr)(status.Code());
              });
        });
    RegisterServiceFunc(
        "{{rpc_func_name}}",
        nullptr,
        aimrt::GetProtobufMessageTypeSupport<{{rpc_req_name}}>(),
        aimrt::GetProtobufMessageTypeSupport<{{rpc_rsp_name}}>(),
        std::move(service_callback));
  }
{{method end}}
}
{{service end}}

{{for service begin}}
{{service_name}}CoService::{{service_name}}CoService() : aimrt::rpc::CoServiceBase(kRpcType, k{{service_name}}Name) {
{{for method begin}}
  {
    aimrt::rpc::ServiceFunc service_callback(
        [this](const aimrt_rpc_context_base_t* ctx, const void* req, void* rsp, aimrt_function_base_t* result_callback_ptr) {
          auto handle_ptr = std::make_unique<const aimrt::rpc::CoRpcHandle>(
              [this](aimrt::rpc::ContextRef ctx_ref, const void* req_ptr, void* rsp_ptr)
                  -> aimrt::co::Task<aimrt::rpc::Status> {
                return {{rpc_func_name}}(
                    ctx_ref,
                    *static_cast<const {{rpc_req_name}}*>(req_ptr),
                    *static_cast<{{rpc_rsp_name}}*>(rsp_ptr));
              });

          aimrt::rpc::ServiceCallback result_callback(result_callback_ptr);

          aimrt::rpc::ContextRef ctx_ref(ctx);

          auto* ptr = handle_ptr.get();
          aimrt::co::StartDetached(
              aimrt::co::On(
                  aimrt::co::InlineScheduler(),
                  filter_mgr_.InvokeRpc(*ptr, ctx_ref, req, rsp)) |
              aimrt::co::Then(
                  [handle_ptr{std::move(handle_ptr)}, result_callback{std::move(result_callback)}](aimrt::rpc::Status status) {
                    result_callback(status.Code());
                  }));
        });
    RegisterServiceFunc(
        "{{rpc_func_name}}",
        nullptr,
        aimrt::GetProtobufMessageTypeSupport<{{rpc_req_name}}>(),
        aimrt::GetProtobufMessageTypeSupport<{{rpc_rsp_name}}>(),
        std::move(service_callback));
  }
{{method end}}
}
{{service end}}

{{for service begin}}
bool Register{{service_name}}ClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name) {
{{for method begin}}
  if (!(rpc_handle_ref.RegisterClientFunc(
          kRpcType,
          service_name,
          "{{rpc_func_name}}",
          nullptr,
          aimrt::GetProtobufMessageTypeSupport<{{rpc_req_name}}>(),
          aimrt::GetProtobufMessageTypeSupport<{{rpc_rsp_name}}>())))
    return false;
{{method end}}
  return true;
}
bool Register{{service_name}}ClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref) {
  return Register{{service_name}}ClientFunc(rpc_handle_ref, k{{service_name}}Name);
}
{{service end}}

{{for service begin}}
{{service_name}}SyncProxy::{{service_name}}SyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref)
    : aimrt::rpc::SyncProxyBase(rpc_handle_ref, kRpcType, k{{service_name}}Name) {}

{{service_name}}SyncProxy::{{service_name}}SyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name)
    : aimrt::rpc::SyncProxyBase(rpc_handle_ref, kRpcType, service_name) {}

{{for method begin}}
aimrt::rpc::Status {{service_name}}SyncProxy::{{rpc_func_name}}(
    aimrt::rpc::ContextRef ctx_ref,
    const {{rpc_req_name}}& req,
    {{rpc_rsp_name}}& rsp) {
  const std::string& full_func_name = aimrt::rpc::GetFullFuncName(rpc_type_, service_name_, "{{rpc_func_name}}");
  return Invoke(full_func_name, ctx_ref, req, rsp);
}
{{method end}}
{{service end}}

{{for service begin}}
{{service_name}}AsyncProxy::{{service_name}}AsyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref)
    : aimrt::rpc::AsyncProxyBase(rpc_handle_ref, kRpcType, k{{service_name}}Name) {}

{{service_name}}AsyncProxy::{{service_name}}AsyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name)
    : aimrt::rpc::AsyncProxyBase(rpc_handle_ref, kRpcType, service_name) {}

{{for method begin}}
void {{service_name}}AsyncProxy::{{rpc_func_name}}(
    aimrt::rpc::ContextRef ctx_ref,
    const {{rpc_req_name}}& req,
    {{rpc_rsp_name}}& rsp,
    std::function<void(aimrt::rpc::Status)>&& callback) {
  const std::string& full_func_name = aimrt::rpc::GetFullFuncName(rpc_type_, service_name_, "{{rpc_func_name}}");
  Invoke(full_func_name, ctx_ref, req, rsp, std::move(callback));
}
{{method end}}
{{service end}}

{{for service begin}}
{{service_name}}FutureProxy::{{service_name}}FutureProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref)
    : aimrt::rpc::FutureProxyBase(rpc_handle_ref, kRpcType, k{{service_name}}Name) {}

{{service_name}}FutureProxy::{{service_name}}FutureProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name)
    : aimrt::rpc::FutureProxyBase(rpc_handle_ref, kRpcType, service_name) {}

{{for method begin}}
std::future<aimrt::rpc::Status> {{service_name}}FutureProxy::{{rpc_func_name}}(
    aimrt::rpc::ContextRef ctx_ref,
    const {{rpc_req_name}}& req,
    {{rpc_rsp_name}}& rsp) {
  const std::string& full_func_name = aimrt::rpc::GetFullFuncName(rpc_type_, service_name_, "{{rpc_func_name}}");
  return Invoke(full_func_name, ctx_ref, req, rsp);
}
{{method end}}
{{service end}}

{{for service begin}}
{{service_name}}CoProxy::{{service_name}}CoProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref)
    : aimrt::rpc::CoProxyBase(rpc_handle_ref, kRpcType, k{{service_name}}Name) {}

{{service_name}}CoProxy::{{service_name}}CoProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name)
    : aimrt::rpc::CoProxyBase(rpc_handle_ref, kRpcType, service_name) {}

{{for method begin}}
aimrt::co::Task<aimrt::rpc::Status> {{service_name}}CoProxy::{{rpc_func_name}}(
    aimrt::rpc::ContextRef ctx_ref,
    const {{rpc_req_name}}& req,
    {{rpc_rsp_name}}& rsp) {
  const std::string& full_func_name = aimrt::rpc::GetFullFuncName(rpc_type_, service_name_, "{{rpc_func_name}}");
  co_return co_await Invoke(full_func_name, ctx_ref, req, rsp);
}
{{method end}}
{{service end}}

{{for service begin}}
void {{service_name}}CoClient::Init(const std::shared_ptr<::aimrt::context::Context> & ctx_ptr, std::string service_name) {
  const ::aimrt::rpc::RpcHandleRef rpc_handle = ::aimrt::context::Context::GetRawRef(*ctx_ptr).GetRpcHandle();

  if (not service_name.empty())
    Register{{service_name}}ClientFunc(rpc_handle, service_name);
  else {
    Register{{service_name}}ClientFunc(rpc_handle);
    service_name = "{{package_name}}.{{service_name}}";
  }

  auto proxy = std::make_shared<{{service_name}}CoProxy>(rpc_handle);
  proxy->SetServiceName(service_name);

{{for method begin}}
  this->{{rpc_func_name}} = ctx_ptr->cli().Init<{{rpc_req_name}}, {{rpc_rsp_name}}>(
    kRpcType.data() + std::string(":/") + service_name + "/{{rpc_func_name}}",
    [proxy](aimrt::rpc::ContextRef ctx_ref, const {{rpc_req_name}}& req, {{rpc_rsp_name}}& rsp)
      -> ::aimrt::co::Task<::aimrt::rpc::Status>
    {
      co_return co_await proxy->{{rpc_func_name}}(ctx_ref, req, rsp);
    }
  );

{{method end}}
}

std::vector<std::string> {{service_name}}CoClient::GetMethodNames(std::string service_name) {
  if (service_name.empty())
    service_name = "{{package_name}}.{{service_name}}";

  return {
{{for method begin}}
    std::string(kRpcType) + ":/" + service_name + "/{{rpc_func_name}}",
{{method end}}
  };
}
{{service end}}

{{for service begin}}
void {{service_name}}CoServer::Init(const std::shared_ptr<::aimrt::context::Context> & ctx_ptr, std::string service_name) {
  if (service_name.empty())
    service_name = "{{package_name}}.{{service_name}}";

  ctx_weak_ptr_ = ctx_ptr;

  impl_ = std::make_shared<{{service_name}}AimRTImpl>(
      std::weak_ptr<::aimrt::context::Context>(ctx_ptr), shared_from_this());

{{for method begin}}
  this->{{rpc_func_name}} = ctx_ptr->srv().Init<{{rpc_req_name}}, {{rpc_rsp_name}}>(
    kRpcType.data() + std::string(":/") + service_name + "/{{rpc_func_name}}"
  );

{{method end}}

  ::aimrt::context::Context::GetRawRef(*ctx_ptr).GetRpcHandle().RegisterService(service_name, impl_.get());
}

std::vector<std::string> {{service_name}}CoServer::GetMethodNames(std::string service_name) {
  if (service_name.empty())
    service_name = "{{package_name}}.{{service_name}}";

  return {
{{for method begin}}
    std::string(kRpcType) + ":/" + service_name + "/{{rpc_func_name}}",
{{method end}}
  };
}

{{service_name}}AimRTImpl::{{service_name}}AimRTImpl(
    std::weak_ptr<::aimrt::context::Context> weak_ctx_ptr,
    std::shared_ptr<{{service_name}}CoServer> res)
    : weak_ctx_ptr_(std::move(weak_ctx_ptr)), res_(std::move(res)) {}

{{for method begin}}
aimrt::co::Task<aimrt::rpc::Status> {{service_name}}AimRTImpl::{{rpc_func_name}}(
    aimrt::rpc::ContextRef ctx_ref,
    const {{rpc_req_name}}& req,
    {{rpc_rsp_name}}& rsp) {
  auto ctx_ptr = weak_ctx_ptr_.lock();

  ctx_ptr->LetMe();
  co_return co_await ctx_ptr->srv().Serving(res_->{{rpc_func_name}}, ctx_ref, req, rsp);
}

{{method end}}
{{service end}}

{{namespace_end}}
)str";

  struct MethodNode {
    std::unordered_map<std::string, std::string> kv;
  };

  struct ServiceNode {
    std::unordered_map<std::string, std::string> kv;
    std::vector<MethodNode> method_vec;
  };

  struct PackageNode {
    std::unordered_map<std::string, std::string> kv;
    std::vector<ServiceNode> service_vec;
  };

  static std::string GenMethodCode(std::string_view temp, const MethodNode& method_node) {
    std::string result(temp);

    for (const auto& item : method_node.kv) {
      ReplaceString(result, item.first, item.second);
    }

    return result;
  }

  static std::string GenServiceCode(std::string_view temp, const ServiceNode& service_node) {
    std::string result(temp);

    for (const auto& item : service_node.kv) {
      ReplaceString(result, item.first, item.second);
    }

    static const std::string kMethodBeginFlag = "{{for method begin}}\n";
    static const std::string kMethodEndFlag = "{{method end}}\n";

    size_t cur_pos = 0;
    while (true) {
      auto begin_pos = result.find(kMethodBeginFlag, cur_pos);
      if (begin_pos == std::string::npos) break;

      auto end_pos = result.find(kMethodEndFlag, begin_pos + kMethodBeginFlag.size());
      if (end_pos == std::string::npos) break;

      std::string cur_temp = result.substr(
          begin_pos + kMethodBeginFlag.size(),
          end_pos - begin_pos - kMethodBeginFlag.size());
      std::string cur_result;
      for (const auto& node : service_node.method_vec) {
        cur_result += GenMethodCode(cur_temp, node);
      }

      result.replace(begin_pos, end_pos - begin_pos + kMethodEndFlag.size(), cur_result);
      cur_pos = begin_pos + cur_result.size();
    }

    return result;
  }

  static std::string GenPackageCode(std::string_view temp, const PackageNode& package_node) {
    std::string result(temp);

    for (const auto& item : package_node.kv) {
      ReplaceString(result, item.first, item.second);
    }

    static const std::string kServiceBeginFlag = "{{for service begin}}\n";
    static const std::string kServiceEndFlag = "{{service end}}\n";

    size_t cur_pos = 0;
    while (true) {
      auto begin_pos = result.find(kServiceBeginFlag, cur_pos);
      if (begin_pos == std::string::npos) break;

      auto end_pos = result.find(kServiceEndFlag, begin_pos + kServiceBeginFlag.size());
      if (end_pos == std::string::npos) break;

      std::string cur_temp = result.substr(
          begin_pos + kServiceBeginFlag.size(),
          end_pos - begin_pos - kServiceBeginFlag.size());
      std::string cur_result;
      for (const auto& node : package_node.service_vec) {
        cur_result += GenServiceCode(cur_temp, node);
      }

      result.replace(begin_pos, end_pos - begin_pos + kServiceEndFlag.size(), cur_result);
      cur_pos = begin_pos + cur_result.size();
    }

    return result;
  }

  static std::string ProtoFileBaseName(const std::string& full_name) {
    return full_name.substr(0, full_name.rfind('.'));
  }

  static std::string GenNamespaceStr(const std::string& ns) {
    std::string result = ns;
    return ReplaceString(result, ".", "::");
  }

  static std::string GenNamespaceBeginStr(const std::string& ns) {
    std::vector<std::string> namespace_vec = SplitToVec(ns, ".");
    std::string result;
    for (const auto& itr : namespace_vec) {
      result += ("namespace " + itr + " {\n");
    }
    return result;
  }

  static std::string GenNamespaceEndStr(const std::string& ns) {
    std::vector<std::string> namespace_vec = SplitToVec(ns, ".");
    std::reverse(namespace_vec.begin(), namespace_vec.end());
    std::string result;
    for (const auto& itr : namespace_vec) {
      result += ("}  // namespace " + itr + "\n");
    }
    return result;
  }

  static void WriteToFile(google::protobuf::compiler::GeneratorContext* context,
                          const std::string& file_name,
                          const std::string& file_context) {
    std::unique_ptr<google::protobuf::io::ZeroCopyOutputStream> output(context->Open(file_name));
    google::protobuf::io::CodedOutputStream coded_out(output.get());
    coded_out.WriteRaw(file_context.data(), file_context.size());
  }

  bool Generate(const google::protobuf::FileDescriptor* file,
                const std::string& parameter,
                google::protobuf::compiler::GeneratorContext* context,
                std::string* error) const override {
    const std::string& file_name = ProtoFileBaseName(file->name());

    PackageNode package_node;
    package_node.kv["{{cur_gencode_version}}"] = "1000000";
    package_node.kv["{{cur_gencode_version_str}}"] = "1.0.0";
    package_node.kv["{{file_name}}"] = file_name;
    package_node.kv["{{package_name}}"] = file->package();
    package_node.kv["{{namespace_begin}}"] = GenNamespaceBeginStr(file->package());
    package_node.kv["{{namespace_end}}"] = GenNamespaceEndStr(file->package());

    for (int ii = 0; ii < file->service_count(); ++ii) {
      const auto* service = file->service(ii);

      ServiceNode service_node;
      service_node.kv["{{service_name}}"] = service->name();

      for (int jj = 0; jj < service->method_count(); ++jj) {
        const auto* method = service->method(jj);

        MethodNode method_node;
        method_node.kv["{{rpc_func_name}}"] = method->name();
        method_node.kv["{{rpc_req_name}}"] = "::" + GenNamespaceStr(method->input_type()->full_name());
        method_node.kv["{{rpc_rsp_name}}"] = "::" + GenNamespaceStr(method->output_type()->full_name());

        service_node.method_vec.emplace_back(std::move(method_node));
      }

      package_node.service_vec.emplace_back(std::move(service_node));
    }

    // hfile
    std::string hfile = GenPackageCode(kTHfile, package_node);
    WriteToFile(context, file_name + ".aimrt_rpc.pb.h", hfile);

    // ccfile
    std::string ccfile = GenPackageCode(kTCcfile, package_node);
    WriteToFile(context, file_name + ".aimrt_rpc.pb.cc", ccfile);

    return true;
  }
};

int main(int argc, char* argv[]) {
  AimRTCodeGenerator generator;
  return google::protobuf::compiler::PluginMain(argc, argv, &generator);
}
